using System.Collections;
using System.Collections.Generic;
using UnityEngine;

namespace JufGame
{
    /// <summary>
    /// 激活函数层，具体激活函数取决于传入的compute shader
    /// </summary>
	public class ActivateFunc : BaseLayer
	{
        private bool isLoadLastOutput;
        
        /// <summary>
        /// 激活函数层初始化
        /// </summary>
        /// <param name="shader">激活函数</param>
        /// <param name="batch">输入批次</param>
        /// <param name="inputCount">每批输入的个数</param>
        /// <param name="isLoadLastOutput">反向传播时是否载入上次输出结果（如sigmoid函数就需要）</param>
        public ActivateFunc(ComputeShader shader, int batch, int inputCount, bool isLoadLastOutput) : base(shader, batch, inputCount)
        {
            this.isLoadLastOutput = isLoadLastOutput;
            outputBuffer = new ComputeBuffer(batch * inputCount, sizeof(float));
        }

        public override void Forward(ComputeBuffer input)
        {
			compute.SetInt("batch", batch);
			compute.SetInt("inputCount", inputCount);
            base.Forward(input);
        }

        public override void Backward(ComputeBuffer dOutputBuffer)
		{
			compute.SetInt("batch", batch);
			compute.SetInt("inputCount", inputCount);
            if(isLoadLastOutput)
            {
                compute.SetBuffer(kernelBackward, "outputBuffer", outputBuffer);
            }
			base.Backward(dOutputBuffer);
		}
	}
}